//
// Created by LUOJIARUI on 2021/4/26.
//

#include "HNSW_Louvain.h"
double precision = 0.000001;
int k = 10;
int thread_num = 4;
int chunk_size = 32;
HNSW_Louvain::HNSW_Louvain() = default;
HNSW_Louvain::HNSW_Louvain(float* data,int d,int nd,long m,vector<long> degree, vector<int> links, vector<float> weight) {
    this->data=data;
    this->nd=nd;this->d=d;
    this->links=links;
    this->weight=weight;
    this->degree=degree;
    Graph g((int) nd, m, (double) m, degree, links, weight);
    vector<vector<int>> sccs;
    scc(g, sccs);
    int threads = 4;
    thread_num = threads < omp_get_max_threads() ? threads : omp_get_max_threads();
    initmaps(g.nb_nodes, thread_num);
    int maxsize = 0;
    int idx = -1;
    for (int i = 0; i < sccs.size(); ++i) {
        idx = i;
        Graph subg(g, sccs[idx]);
        Community c(subg, -1, precision, thread_num, chunk_size);
        // the hierarchy structure for the largest component
        //需要记录每个community中心元素的id
        vector<vector<pair<int, vector<int>>>> partit;
        vector<vector<pair<int,int>>> up2bot;
        vector<vector<int>> lss;
        vector<vector<long>> dss;
        vector<vector<float>> wss;
        louvainPlus(c, partit,lss,wss,dss,up2bot);
        if (!partit.empty()) {
            partitions.emplace_back(partit);
            up2bottom.emplace_back(up2bot);
            ls.push_back(lss);
            ws.push_back(wss);
            ds.push_back(dss);
        }
        else {
/*
            vector<pair<int,vector<int>>> sccc;
            for (int ji=0;ji<sccs[i].size();ji++){
                pair<int,vector<int>> pp(sccs[i][0],sccs[i]);
                sccc.push_back(pp);
            }
            partit.push_back(sccc);
            partitions.emplace_back(partit);
*/
        }
    }
}

float
HNSW_Louvain::dis(const float *data1,const float *data2) {
    float nowDis=0;
    for (int i=0;i<d;i++){
        nowDis+=pow(*(data1+i)-*(data2+i),2);
    }
    return nowDis;
}
using namespace __gnu_cxx;
myPair
HNSW_Louvain::searchUpperLayer(int selectSCC,int layerNum, myPair epPair,float* queryPos) {
    int ep=epPair.first;
    float epDis=0;
    hash_set<int> v;
    int id=up2bottom[selectSCC][layerNum][ep].first;
    v.insert(ep);
    priority_queue<myPair,vector<myPair>,ComparatorForMinHeap> C;
    epDis=dis(queryPos,data+id*d);
    myPair p(ep,epDis);
    C.push(p);
    while (!C.empty()){
        myPair pp=C.top();
        C.pop();
        if (pp.second>epDis){
            break;
        }
        int sub=pp.first;
        long start=(sub==0)?0:ds[selectSCC][layerNum][sub-1];
        long end=ds[selectSCC][layerNum][sub];
        for (long jj=start;jj<end;jj++){
            int index=ls[selectSCC][layerNum][jj];
            if (v.count(index)==0){
                v.insert(index);
                int indID=up2bottom[selectSCC][layerNum][index].first;
                float nowDis=0;
                nowDis=dis(queryPos,data+indID*d);
                if (nowDis<epDis){
                    myPair nowPair(index,nowDis);
                    C.push(nowPair);
                    epDis=nowDis;
                    ep=index;
                }
            }
        }
    }
    ep=(layerNum!=0)?up2bottom[selectSCC][layerNum][ep].second:up2bottom[selectSCC][0][ep].first;
    myPair res(ep,epDis);
    return res;
}

priority_queue<myPair,vector<myPair>,ComparatorForMaxHeap>
HNSW_Louvain::searchBaseLayer(int k1,int ef,myPair epPair,float* queryPos){
    int ep=epPair.first;
    float epDis=epPair.second;
    hash_set<int> v;
    v.insert(ep);
    priority_queue<myPair,vector<myPair>,ComparatorForMaxHeap> W;
    priority_queue<myPair,vector<myPair>,ComparatorForMinHeap> C;
    myPair pd(ep,epDis);
    W.push(pd);
    C.push(pd);
    while (!C.empty()){
        myPair minPair=C.top();
        C.pop();
        myPair maxPair=W.top();
        if (minPair.second>maxPair.second){
            break;
        }
        int ig=minPair.first;
        long st=(ig==0)?0:degree[ig-1];
        long en=degree[ig];
        for (long t=st;t<en;t++) {
            int index = links[t];
            if (v.count(index) == 0) {
                v.insert(index);
                float nowDis = dis(queryPos,data+index*d);
                maxPair = W.top();
                if (nowDis < maxPair.second || W.size() < ef) {
                    myPair nowPair(index, nowDis);
                    W.push(nowPair);
                    C.push(nowPair);
                    if (W.size() > ef) {
                        W.pop();
                    }
                }
            }
        }
    }
    while (W.size()>k1){
        W.pop();
    }
    return W;
}

priority_queue<myPair,vector<myPair>,ComparatorForMaxHeap>
HNSW_Louvain::knn(float* queryPos,int k1,int ef){
    int ep=0;
    float epDis=0;
    int selectSCC=0;
    int maxLayer=up2bottom[selectSCC].size();
    int bsfID=up2bottom[selectSCC][maxLayer-1][ep].first;
    float bsfSCCDis=dis(queryPos,data+bsfID*d);
    if (up2bottom.size()>1){
        for (int j=1;j<up2bottom.size();j++){
            maxLayer=up2bottom[j].size();
            int nowID=up2bottom[j][maxLayer-1][ep].first;
            float nowSCCDIS=dis(queryPos,data+nowID*d);
            if (nowSCCDIS<bsfSCCDis){
                bsfSCCDis=nowSCCDIS;
                selectSCC=j;
            }
        }
    }
    maxLayer=up2bottom[selectSCC].size();
    myPair epPair(ep,epDis);
    for (int j=maxLayer-1;j>=0;j--){
        epPair=searchUpperLayer(selectSCC,j,epPair,queryPos);
    }
    return searchBaseLayer(k1,ef,epPair,queryPos);
}

vector<vector<int>> &scc(Graph &G, vector<vector<int>> &SCCG) {
    int vnums = G.nb_nodes;
    vector<int> DFN(vnums, 0), q(vnums, 0), LOW(vnums, 0), st(vnums, 0), viter(vnums, 0), id(vnums, 0);
    vector<bool> visit(vnums, false);
    int top, Tindex, stop, snum = 0;
    top = -1;
    Tindex = 0;
    stop = -1;
    for (int u = 0; u < vnums; u++) {
        if (!DFN[u]) {
            int v;
            DFN[u] = LOW[u] = ++Tindex;
            visit[u] = true;
            q[++top] = u;
            st[++stop] = u;
            while (stop != -1) {
                u = st[stop];
                pair<vector<int>::iterator, vector<float>::iterator> p = G.neighbors(u);
                for (int iter = viter[u]; iter < G.nb_neighbors(u); iter++) {
                    v = *(p.first + iter);
                    if (!DFN[v]) {
                        DFN[v] = LOW[v] = ++Tindex;
                        visit[v] = true;
                        q[++top] = v;
                        st[++stop] = v;
                        break;
                    } else if (DFN[v] > DFN[u] && LOW[v] < LOW[u]) {
                        LOW[u] = LOW[v];
                    } else if (visit[v] && DFN[v] < LOW[u])
                        LOW[u] = DFN[v];
                }
                if (u == st[stop]) {
                    if (DFN[u] == LOW[u]) {
                        do {
                            v = q[top--];
                            visit[v] = false;
                            id[v] = snum;
                        } while (v != u);
                        snum++;
                    }
                    stop--;
                }
            }
        }
    }
    SCCG.resize(snum);
    DFN.assign(vnums, 0);
    for (int i = 0; i < vnums; i++) {
        SCCG[id[i]].emplace_back(i);
    }
    return SCCG;
}

vector<vector<pair<int, vector<int>>>> &louvainPlus(Community &c, vector<vector<pair<int, vector<int>>>> &partit,
                                                    vector<vector<int>> &linksInLayer,
                                                    vector<vector<float>> &weightInLayer,
                                                    vector<vector<long>> &degreeInLayer,
                                                    vector<vector<pair<int,int>>>& up2bottom) {
    Graph g;
    double mod = c.modularity(), new_mod;
    int level = 0;

// shiqi: louvain+ phase tries to merge communities as many as possible
    while (c.g.nb_nodes > k) {
        c.one_level_new(k);
        new_mod = c.modularity();
        vector<pair<int, vector<int>>> comm_nodes;
        vector<pair<int,int>> u2b;
// shiqi: generate new compound graph and store into partitions
        g = c.partition2graph_binary(comm_nodes,u2b);
        up2bottom.push_back(u2b);
        degreeInLayer.push_back(g.degrees);
        linksInLayer.push_back(g.links);
        weightInLayer.push_back(g.weights);
//        g = c.partition2graph_binary_old(comm_nodes);
        partit.push_back(comm_nodes);
        c = Community(g, -1, precision, thread_num, chunk_size);
        mod = new_mod;
//        cout<<"level time:"<<time_by(start)<<endl;
//        if (verbose)
//            display_time("  end computation");

    }
    return partit;
}